#include "imagemask.h"
#include <set>
#include <opencv2/highgui/highgui.hpp>
#include <random>
#include <iostream>


using namespace cv;

ImageMask::ImageMask()
    :mNumberOfLayers(2), mNumComponents(-1)
{
}

ImageMask::ImageMask(const ImageMask& inputMask)
    :mNumberOfLayers(inputMask.mNumberOfLayers), mNumComponents(-1)
{
    mMask=inputMask.mMask.clone();
}

ImageMask::ImageMask(const Size& size,unsigned int numberOfLayers)
    :ImageMask()
{
    create(size,numberOfLayers);
}

void ImageMask::create(const Size& size,unsigned int numberOfLayers)
{
    mMask.create(size,CV_8UC1);
    mNumberOfLayers=numberOfLayers;
}

void ImageMask::connectedComponents() {
    int nrows = mMask.rows;
    int ncols = mMask.cols;
    cv::Mat labels(mMask.size(), CV_8UC1);
    labels.setTo(Scalar(0));
    std::vector<unsigned int> ufLabels;
    ufLabels.push_back(0);
    
    // first pass
    int currentLabel = 1;
    std::set<unsigned int> neighborLabels;
    for (int x = 0; x < nrows; x++) {
        for (int y = 0; y < ncols; y++) {
            if (mMask.at<uchar>(x, y) == 0) continue;
            // find neighbor labels 
            neighborLabels.clear();
            if(x>0) {  // top
                if (mMask.at<uchar>(x-1, y) != 0) 
                    neighborLabels.insert(labels.at<uchar>(x-1, y));}
            if(y>0) {  // left
                if (mMask.at<uchar>(x, y-1) != 0) 
                    neighborLabels.insert(labels.at<uchar>(x, y-1));}
            if(x>0 && y>0) {  // top-left
                if (mMask.at<uchar>(x-1, y-1) != 0) 
                    neighborLabels.insert(labels.at<uchar>(x-1, y-1));}
            if(x>0 && y<ncols-1) { // top-right
                if (mMask.at<uchar>(x-1, y+1) != 0) 
                    neighborLabels.insert(labels.at<uchar>(x-1, y+1));}
            // if neighbors empty use new label
            if(neighborLabels.empty()) {
                labels.at<uchar>(x, y) = currentLabel;
                ufLabels.push_back(currentLabel);
                currentLabel++;}
            // else assign to smallest neighbor and do union on neighbors
            else {
                unsigned int minLabel = *std::min_element(std::begin(neighborLabels), std::end(neighborLabels));
                labels.at<uchar>(x, y) = minLabel;
                std::set<unsigned int>::iterator it;
                if (neighborLabels.size() == 1) continue;
                for (it=neighborLabels.begin(); it!=neighborLabels.end(); it++) 
                    ufUnion(*it, minLabel, ufLabels);}
        }
    }
    // second pass 
    std::set<unsigned int> unique_labels;
    for (int x = 0; x < nrows; x++) {
        for (int y = 0; y < ncols; y++) {
            if (labels.at<uchar>(x, y) == 0) continue;  
            unsigned int label = ufFind(labels.at<uchar>(x, y), ufLabels);
            labels.at<uchar>(x, y) = label;
            unique_labels.insert(label);
        }
    }
    // relabel components
    mComponentLabels.create(labels.size(), CV_8UC1);
    mComponentLabels.setTo(Scalar(0));
    unsigned int index = 1;
    for (std::set<unsigned int>::iterator it=unique_labels.begin(); it!=unique_labels.end(); ++it) {
        mComponentLabels.setTo(Scalar(index), labels == *it);
        index++;
    }
    // number of components
    mNumComponents = (int)unique_labels.size();
}

int ImageMask::largestComponent() {
    assert(mNumComponents != -1);
    unsigned int max_index = 0;
    double max_area = 0.0;
    for (int cc = 1; cc <= mNumComponents; cc++) {
        double area = sum(mComponentLabels == cc).val[0];
        if (area > max_area) { 
            max_area = area;
            max_index = cc;
        }
    }
    return (int)max_index;
}

void ImageMask::visualizeConnectedComponents() {
    assert(mNumComponents != -1);
    cv::Mat visualizer(mMask.size(), CV_8UC3);
    RNG rng(12345);
    visualizer.setTo(Scalar(0, 0, 0));
    for (unsigned int cc = 1; cc <= (unsigned int)mNumComponents; cc++) {
        Scalar color = Scalar(rng.uniform(0,255), rng.uniform(0, 255), rng.uniform(0, 255));
        visualizer.setTo(color, mComponentLabels == cc);
    }
    namedWindow( "Connected Components", CV_WINDOW_AUTOSIZE );
    imshow( "Connected Components", visualizer);
    waitKey(0);
}

void ImageMask::componentBoundary() {
    assert(mNumComponents != -1);
    for (unsigned int cc = 0; cc <= (unsigned int)mNumComponents; cc++) 
        mComponentBoundary.push_back(std::vector<Point>());
    int nrows = mMask.rows;
    int ncols = mMask.cols;
    for (int x = 0; x < nrows; x++) {
        for (int y = 0; y < ncols; y++) {
            unsigned int cc = mComponentLabels.at<uchar>(x, y);
            if (cc == 0) continue;
            if (x==0 || x==nrows-1 || y==0 || y==ncols-1) {
                mComponentBoundary[cc].push_back(Point(y, x));
                continue;}
            bool isBoundary = false;
            for (int i = x-1; i <= x+1; i++) {
                for (int j = y-1; j <= y+1; j++) {
                    //if (i<0 || i>=nrows || j<0 || j>=ncols) continue;
                    if(mComponentLabels.at<uchar>(i, j) == 0) isBoundary = true;
                }
            }
            if (isBoundary) 
                mComponentBoundary[cc].push_back(Point(y, x));
        }
    }
}

void ImageMask::boundary() {
    int nrows = mMask.rows;
    int ncols = mMask.cols;
    for (int x = 0; x < nrows; x++) {
        for (int y = 0; y < ncols; y++) {
            unsigned int cc = mMask.at<uchar>(x, y);
            if (cc == 0) continue;
            bool isBoundary = false;
            for (int i = x-1; i <= x+1; i++) {
                for (int j = y-1; j <= y+1; j++) {
                    if (i<0 || i>=nrows || j<0 || j>=ncols) continue;
                    if(mMask.at<uchar>(i, j) == 0) isBoundary = true;
                }
            }
            if (isBoundary)
                mBoundary.push_back(Point(y, x));
        }
    }
}

void ImageMask::visualizeComponentsBoundary() {
    assert(mNumComponents != -1);
    assert(mComponentBoundary.size() != 0);
    cv::Mat visualizer(mMask.size(), CV_8UC3);
    RNG rng(12345);
    visualizer.setTo(Scalar(0, 0, 0));
    for (unsigned int cc = 1; cc <= (unsigned int)mNumComponents; cc++) {
        Vec3b color(rng.uniform(0,255), rng.uniform(0, 255), rng.uniform(0, 255));
        std::vector<Point> contour = mComponentBoundary[cc];
        for (int i = 0; i < (int)contour.size(); i++)
            visualizer.at<Vec3b>(contour[i].y, contour[i].x) = color;
    }
    namedWindow( "Connected Components Boundary", CV_WINDOW_AUTOSIZE );
    imshow( "Connected Components Boundary", visualizer);
    waitKey(0);
}


Mat ImageMask::applyTo(const std::vector<const cv::Mat*> images)
{
    assert(mNumberOfLayers>0 && images.size()==mNumberOfLayers);
    Mat output(images[0]->size(), CV_8UC3);
    output.setTo(Vec3b(255,255,255));
    for(unsigned int layer=0;layer<mNumberOfLayers;layer++)
    {
        images[layer]->copyTo(output,mMask==layer);
    }
    return output;
}


Mat ImageMask::applyTo(const Mat* image) {
    Mat output(image->size(), CV_8UC3);
    output.setTo(Vec3b(0,0,0));
    image->copyTo(output,mMask==1);
    return output; 
}


unsigned int ImageMask::ufFind(unsigned int x, std::vector<unsigned int>& unionFind) {
    assert(x >= 0 && x < unionFind.size());
    if (unionFind[x] == x) return x;
    else return ufFind(unionFind[x], unionFind);
}

void ImageMask::ufUnion(unsigned int x, unsigned int y, std::vector<unsigned int>& unionFind) {
    unsigned int xParent = ufFind(x, unionFind);
    unsigned int yParent = ufFind(y, unionFind);
    unionFind[xParent] = yParent;
}

